
import asyncio
import aiohttp
from openai import AsyncOpenAI

# 初始化模型客户端
client = AsyncOpenAI(
    base_url="http://36.248.221.175:8012/v1",
    api_key="EMPTY"
)

# 天气查询工具 - 这里使用模拟数据，实际应用中可替换为真实天气API
async def get_weather(city: str) -> str:
    """获取指定城市的天气预报"""
    # 模拟天气API调用
    print(f"查询 {city} 的天气信息...")
    
    # 实际应用中可替换为真实API，例如：
    # api_key = "你的天气API密钥"
    # url = f"https://api.weatherapi.com/v1/forecast.json?key={api_key}&q={city}&days=3"
    # async with aiohttp.ClientSession() as session:
    #     async with session.get(url) as response:
    #         data = await response.json()
    #         return parse_weather_data(data)
    
    # 模拟返回的天气数据
    mock_weather = {
        "北京": "晴，22-30°C，西北风3级，空气质量优",
        "上海": "多云转阴，24-28°C，东南风2级，空气质量良好",
        "广州": "雷阵雨，26-32°C，南风4级，空气质量一般",
        "深圳": "多云，25-31°C，东风3级，空气质量良好"
    }
    
    if city in mock_weather:
        return f"{city}今日天气：{mock_weather[city]}"
    else:
        return f"抱歉，暂未获取到{city}的天气信息"

async def weather_agent(user_query: str) -> str:
    """天气预报Agent主函数"""
    # 1. 让模型解析用户查询，提取城市信息
    system_prompt = """
    你是一个天气查询助手。请分析用户的问题，提取需要查询天气的城市名称。
    如果能明确提取城市，直接返回城市名称（仅返回城市名，不要其他内容）。
    如果无法确定城市，返回"需要城市名称"。
    """
    
    parse_response = await client.chat.completions.create(
        model="qwen",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_query}
        ],
        temperature=0.1,  # 降低随机性，提高提取准确性
        max_tokens=50
    )
    
    city = parse_response.choices[0].message.content.strip()
    
    # 2. 处理提取结果
    if city == "需要城市名称":
        return "请问你想查询哪个城市的天气？"
    
    # 3. 调用天气工具获取数据
    weather_data = await get_weather(city)
    
    # 4. 整理并返回结果
    response = await client.chat.completions.create(
        model="qwen",
        messages=[
            {"role": "system", "content": "你是一个专业的天气预报员，请将天气信息整理成自然、友好的回答。"},
            {"role": "user", "content": f"请播报{city}的天气：{weather_data}"}
        ],
        temperature=0.7,
        max_tokens=200
    )
    
    return response.choices[0].message.content

async def main():
    print("天气预报助手已启动，输入'退出'结束对话")
    
    while True:
        user_input = input("请输入你的查询：")
        if user_input.lower() == "退出":
            print("再见！")
            break
            
        result = await weather_agent(user_input)
        print(f"天气助手：{result}\n")

if __name__ == "__main__":
    asyncio.run(main())
    